from datasets import load_dataset
from sentence_transformers import SentenceTransformer, losses, SentenceTransformerTrainer, \
    SentenceTransformerTrainingArguments, SimilarityFunction, SentenceTransformerModelCardData
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator, TripletEvaluator

# finetune bge-m3 pip install 'accelerate>=0.26.0'
# 参考文档： https://www.sbert.net/docs/sentence_transformer/training_overview.html#why-finetune

# Load a model to finetune with
embedding_model = SentenceTransformer(model_name_or_path="../transformers/bge-m3",
                                      local_files_only=True,
                                      model_card_data=SentenceTransformerModelCardData(
                                          language="zh_cn",
                                          license="apache-2.0",
                                          model_name="bgm-m3-zh",
                                      ))
# 2.加载数据集
dataset = load_dataset(path="../data", data_files={"train": "train_data.json",
                                                   "eval": "eval_data.json",
                                                   "test": "test_data.json"}, )
train_data = dataset.get("train")
eval_data = dataset.get("eval")
test_dataset = dataset.get("test")

# 训练参数
args = SentenceTransformerTrainingArguments(
    output_dir="../checkpoints",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    warmup_steps=100
)

# 3.定义loss function
losses_function = losses.MultipleNegativesRankingLoss(model=embedding_model)

# 4 定义结果集评估
evaluator = TripletEvaluator(
    anchors=[str(q) for q in eval_data["query"]],
    positives=[str(p) for p in eval_data["positive_passage"]],
    negatives=[str(n) for n in eval_data["negative_passage"]],
    main_similarity_function=SimilarityFunction.COSINE,
    name="bge-m3-eval"
)

trainer = SentenceTransformerTrainer(
    model=embedding_model,
    args=args,
    loss=losses_function,
    train_dataset=train_data,
    eval_dataset=eval_data,
    evaluator=evaluator
)

# 4.开始训练
trainer.train()

# 完成之后，进行评估
test_evaluator = TripletEvaluator(
    anchors=[str(q) for q in test_dataset["query"]],
    positives=[str(p) for p in test_dataset["positive_passage"]],
    negatives=[str(n) for n in test_dataset["negative_passage"]],
    main_similarity_function=SimilarityFunction.COSINE,
    name="bge-m3-test",
    show_progress_bar=True
)

# 使用该模型进行评估
result = embedding_model.evaluate(test_evaluator, output_path="../output")
print(f"Test result: {result}")

# 保存训练模型
embedding_model.save_pretrained(path="../transformers/fin-bge-m3")
